package com.ctfo.ulcp.exp.db

import com.ctfo.ulcp.exp.core.CrudEmbeddingStore
import com.ctfo.ulcp.exp.task.LongRunningTaskManager
import com.ctfo.ulcp.exp.utils.JdbcService
import com.ctfo.ulcp.exp.utils.JsonCodec
import com.ctfo.ulcp.exp.utils.MySqlService
import com.ctfo.ulcp.exp.utils.TableInfo
import dev.langchain4j.data.document.Metadata
import dev.langchain4j.data.document.splitter.DocumentSplitters
import dev.langchain4j.data.segment.TextSegment
import dev.langchain4j.model.embedding.EmbeddingModel
import org.jetbrains.exposed.sql.*
import org.slf4j.LoggerFactory

class ResourceManager(private val embeddingModel: EmbeddingModel, private val embeddingStore: CrudEmbeddingStore) {

    private val logger = LoggerFactory.getLogger(Tasks::class.java)

    fun execSaveResourceInfo(resource: Resource) {
        // 锁定表
        Resources.selectAll().where { Resources.id eq resource.id!! }.forUpdate()
            .map { Resources.transform(it) }
            .firstOrNull()
            ?: throw RuntimeException("resource not exists")

        // 查询是否有运行中的任务
        val pendingTask = Tasks.selectAll()
            .where { (Tasks.status eq TaskStatus.Pending.name).and { Tasks.resourceId eq resource.id!! } }
            .limit(1)
            .firstOrNull()
        if (pendingTask != null) {
            logger.warn("pending task exists for resource:{}", resource)
            return
        }

        // 创建任务
        val taskId = Tasks.createTask(resource.id!!)
        // 提交任务
        LongRunningTaskManager.submitTask(taskId) {
            saveResourceInfo(resource)
        }
    }

    fun dropCollection() {
        // 初始清空数据
        embeddingStore.dropCollection()
    }

    fun searchEmbedding(text: String, resourceId: Int?, resourceType: ResourceType?) {
        val query = Resources.selectAll()
        resourceId?.let { query.andWhere { Resources.id eq resourceId } }
        resourceType?.let { query.andWhere { Resources.type eq resourceType.name } }
        val resources = query.map { Resources.transform(it) }


    }


    fun searchDocEmbedding(text: String, maxResources: Int): String {
        val embedding = embeddingModel.embed(text).content()
        val query = mapOf(Constants.META_KEY_SOURCE_TYPE to EmbeddingSource.SENTENCE.name)
        val matches = embeddingStore.findRelevant(embedding, query, maxResources, 0.0)
        return matches.joinToString("---------------\n") { match -> match.embedded().text() }
    }

    /**
     * 在当前线程执行
     * 是否需要锁定？ 提交任务的时候已经锁定了
     */
    private fun saveResourceInfo(resource: Resource) {
        logger.info("saveResourceInfo:{}", resource)

        addResourceToEmbeddingStore(resource)
        when (resource.property) {
            is InterfaceProperty -> {}//接口没什么要做的
            is DocProperty -> {
                val doc = resource.property.loadDocument()
                doc?.let {
                    val splitter = DocumentSplitters.recursive(300, 0)
                    val segments = splitter.split(doc)
                    addSentenceToEmbeddingStore(resource.id!!, segments)
                }
            }

            is DbProperty -> {
                val property = resource.property
                val jdbcService: JdbcService = when (val dbType = property.dbType) {
                    DataSourceType.MySQL -> MySqlService(
                        property.jdbcUrl,
                        property.user,
                        property.password
                    )

                    else -> throw RuntimeException("not support database: $dbType")
                }

                val tables = jdbcService.use {
                    it.fetchTableInfoListWithSample()
                }

                addTableToEmbeddingStore(resource.id!!, tables)
            }
        }
    }

    private fun addTableToEmbeddingStore(resourceId: Int, tables: List<TableInfo>) {
        DSL.transaction {
            for (table in tables) {
                val exists = TableDescriptors.findBy(resourceId, table.name)
                // 如果已经存在人工处理过的提示词则不再重复处理
                if (exists != null && exists.commentSrc == CommentSrc.HUMAN) {
                    continue
                }

                val info = TableDescriptor(
                    exists?.id,
                    resourceId,
                    table.name,
                    table.comment ?: "",
                    CommentSrc.AUTO
                )
                TableDescriptors.save(info)

                // 保存到向量库
                addTextEmbedding(info.summary, metadataOfTable(resourceId, info.name))
            }
        }
    }


    private fun addSentenceToEmbeddingStore(resourceId: Int, segments: List<TextSegment>) {
        val deleteByCondition = mapOf(
            Constants.META_KEY_RES_ID to resourceId.toString(),
            Constants.META_KEY_SOURCE_TYPE to EmbeddingSource.SENTENCE.name,
        )
        // 对于文档每次保存之前删除同一文档之前已经保存过的数据
        // 其他类型不需要提前删除: 因为id由元数据生成, 和业务数据强关联
        embeddingStore.batchDelete(deleteByCondition)


        segments.forEach { seg ->
            addTextEmbedding(seg.text(), metadataOfSentence(resourceId, seg))
        }
    }

    private fun addResourceToEmbeddingStore(resource: Resource) {
        addTextEmbedding(resource.summary, metadataOfResource(resource))
    }


    private fun addTextEmbedding(text: String, metadata: Map<String, String>) {
        if (text.isNotEmpty()) {
            val segment = TextSegment.from(text, Metadata.from(metadata))
            val embedding = embeddingModel.embed(segment).content()
            embeddingStore.add(embedding, segment)
        }
    }


    private fun metadataOfResource(resource: Resource): Map<String, String> {
        val embeddingSource = when (resource.type) {
            ResourceType.DOC -> EmbeddingSource.DOC
            ResourceType.API -> EmbeddingSource.API
            ResourceType.DB -> EmbeddingSource.DB
        }
        return mapOf(
            Constants.META_KEY_RES_ID to resource.id.toString(),
            Constants.META_KEY_SOURCE_TYPE to embeddingSource.name
        )
    }

    private fun metadataOfTable(resourceId: Int, tableName: String): Map<String, String> {
        return mapOf(
            Constants.META_KEY_RES_ID to resourceId.toString(),
            Constants.META_KEY_SOURCE_TYPE to EmbeddingSource.TABLE.name,
            Constants.META_KEY_SOURCE_ID to tableName
        )
    }

    private fun metadataOfSentence(resourceId: Int, textSegment: TextSegment): Map<String, String> {
        return mapOf(
            Constants.META_KEY_RES_ID to resourceId.toString(),
            Constants.META_KEY_SOURCE_TYPE to EmbeddingSource.SENTENCE.name,
            Constants.META_KEY_SOURCE_ID to textSegment.metadata()["index"]
        ) + textSegment.metadata().asMap()
    }


}


data class Resource(
    var id: Int?,
    var projectId: Int?,
    val type: ResourceType,
    val name: String,
    val summary: String,
    val property: ResourceProperty,
)

object Resources : Table("t_resource") {

    val id = integer("ID").autoIncrement()
    private val projectId = integer("PROJECT_ID")
    val type = varchar("TYPE", length = 255)
    private val name = varchar("NAME", length = 255)
    private val summary = text("SUMMARY")
    private val property = text("PROPERTY")

    override val primaryKey = PrimaryKey(id)

    fun create(resource: Resource): Resource {
        resource.id = Resources.insert {
            it[projectId] = resource.projectId!!
            it[name] = resource.name
            it[type] = resource.type.name
            it[summary] = resource.summary
            it[property] = JsonCodec.encode(resource.property)
        }[id]
        return resource
    }

    fun findByProjectId(projectId: Int): List<Resource> {
        return Resources.selectAll()
            .where { Resources.projectId eq projectId }.map { transform(it) }
    }

    fun transform(row: ResultRow): Resource {
        val resourceType = ResourceType.valueOf(row[type])
        val ps = row[property]
        val property = when (resourceType) {
            ResourceType.API -> JsonCodec.decodeValue(ps, InterfaceProperty::class.java)
            ResourceType.DB -> JsonCodec.decodeValue(ps, DbProperty::class.java)
            ResourceType.DOC -> JsonCodec.decodeValue(ps, DocProperty::class.java)
        }

        return Resource(
            row[id],
            row[projectId],
            resourceType,
            row[name],
            row[summary],
            property
        )
    }


}

fun main(){
    println("1111")
}